# ================================ 3.增量微调 ==============================================
import torch
from transformers import BertModel

# 定义设备信息
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(DEVICE)

# 加载预训练模型
pretrained_model = BertModel.from_pretrained(
    r"D:\code\base_learn\learn_ai\hugging_face\model\bert-base-chinese\models--bert-base-chinese\snapshots\c30a6ed22ab4564dc1e3b2ecbf6e766b0611a33f"
).to(DEVICE)
print(pretrained_model)


# 定义下游任务（增量模型）
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 设计全连接网络，实现二分类任务
        self.fc = torch.nn.Linear(768, 2)

    # 使用模型处理数据（执行前向计算）
    def forward(self, input_ids, attention_mask, token_type_ids):
        # 冻结 Bert 模型的参数，让其不参与训练
        with torch.no_grad():
            outputs = pretrained_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                token_type_ids=token_type_ids
            )
        # # 不冻结的话为全量微调，整个模型都会参与训练
        # outputs = pretrained_model(
        #     input_ids=input_ids,
        #     attention_mask=attention_mask,
        #     token_type_ids=token_type_ids
        # )
        # 增量模型参与训练
        outputs = self.fc(outputs.last_hidden_state[:, 0])
        return outputs
